package com.crscd.ai.store;

import com.crscd.ai.entity.ChatMessages;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ChatMessageDeserializer;
import dev.langchain4j.data.message.ChatMessageSerializer;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import java.util.Collections;
import java.util.List;
import java.util.Objects;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;
import org.springframework.stereotype.Component;

/**
 * Created with IntelliJ IDEA.
 *
 * @description:
 * @author: liuziyang
 * @since: 9/24/25 11:38 AM
 * @modifiedBy:
 * @version: 1.0
 */
@Component
public class MongoChatMemoryStore implements ChatMemoryStore {
  @Autowired private MongoTemplate mongoTemplate;

  /**
   * @param memoryId The ID of the chat memory.
   * @return
   */
  @Override
  public List<ChatMessage> getMessages(Object memoryId) {
    Criteria criteria = Criteria.where("memoryId").is(memoryId);
    Query query = new Query(criteria);
    final ChatMessages one = mongoTemplate.findOne(query, ChatMessages.class);
    if (Objects.isNull(one)) {
      return Collections.emptyList();
    }
    final String contentJson = one.getContent();
    return ChatMessageDeserializer.messagesFromJson(contentJson);
  }

  /**
   * @param memoryId The ID of the chat memory.
   * @param messages List of messages for the specified chat memory, that represent the current
   *     state of the {@link ChatMemory}. Can be serialized to JSON using {@link
   *     ChatMessageSerializer}.
   */
  @Override
  public void updateMessages(Object memoryId, List<ChatMessage> messages) {
    Criteria criteria = Criteria.where("memoryId").is(memoryId);
    Query query = new Query(criteria);
    Update update = new Update();
    update.set("content", ChatMessageSerializer.messagesToJson(messages));
    mongoTemplate.upsert(query, update, ChatMessages.class);
  }

  /**
   * @param memoryId The ID of the chat memory.
   */
  @Override
  public void deleteMessages(Object memoryId) {
    Criteria criteria = Criteria.where("memoryId").is(memoryId);
    Query query = new Query(criteria);
    mongoTemplate.remove(query, ChatMessages.class);
  }
}
